#include <stdio.h>
#include "blas.h"
#include "gemm.h"
#include "im2col.h"
#include "activations.h"
#include "convolutional_layer.h"


int convolutional_out_height(convolutional_layer l) {
    return (l.h + 2 * l.pad - l.size) / l.stride + 1;
}

int convolutional_out_width(convolutional_layer l) {
    return (l.w + 2 * l.pad - l.size) / l.stride + 1;
}

void forward_convolutional_layer(convolutional_layer l, network net)
{
    int i, j;

    fill_cpu(l.outputs * l.batch, 0, l.output, 1);

    int m = l.n / l.groups;
    int k = l.size * l.size * l.c / l.groups;
    int n = l.out_w * l.out_h;
    for (i = 0; i < l.batch; i++) {
        for (j = 0; j < l.groups; j++) {
            float *a = l.weights + j * l.nweights / l.groups;
            float *b = net.workspace;
            float *c = l.output + (i * l.groups + j) * n * m;
            float *im = net.input + ( i * l.groups + j) * l.c / l.groups * l.h * l.w;

            if (l.size == 1) {
                b = im;
            } else {
                im2col_cpu(im, l.c / l.groups, l.h, l.w, l.size, l.stride, l.pad, b);
            }
            gemm(0,0,m,n,k,1,a,k,b,n,1,c,n);
        }
    }

    // if (l.batch_normalize) {

    // } else
    // {
        add_bias(l.output, l.biases, l.batch, l.n, l.out_h * l.out_w);
    // }
    activate_array(l.output, l.outputs * l.batch, l.activation);
}

static size_t get_workspace_size(layer l){
    return (size_t)(l.out_h * l.out_w * l.size * l.size * l.c / l.groups * sizeof(float));
}

convolutional_layer make_convolutional_layer(int batch, int h, int w, int c, 
    int n, int groups, int size, int stride, int padding, 
    ACTIVATION activation, int batch_normalize, int binary, 
    int xnor, int adam)
{
    int i;
    convolutional_layer l = {0};
    l.type = CONVOLUTIONAL;
    l.groups = groups;
    l.h = h;
    l.w = w;
    l.c = c;
    l.n = n;
    l.binary = binary;
    l.xnor = xnor;
    l.batch = batch;
    l.stride = stride;
    l.size = size;
    l.pad = padding;
    l.batch_normalize = batch_normalize;
    l.activation = activation;

    l.nweights = c / groups * n * size * size;
    l.weights = calloc(l.nweights, sizeof(float));

    l.nbiases = n;
    l.biases = calloc(n, sizeof(float));

    l.out_h = convolutional_out_height(l);
    l.out_w = convolutional_out_width(l);
    l.out_c = n;
    l.outputs = l.out_h * l.out_w * l.out_c;
    l.inputs = l.h * l.w * l.c;

    l.output = calloc(l.batch * l.outputs, sizeof(float));
    l.delta =  calloc(l.batch * l.outputs, sizeof(float));

    l.workspace_size = get_workspace_size(l);
    l.forward = forward_convolutional_layer;

    return l;
}